import os

MOT_set = 'MOT15'
if MOT_set == 'MOT17Det':
    root = '/home/sean/data/MOT/MOT15Det/train'
    train_video_list = ['MOT17-02', 'MOT17-04', 'MOT17-05', 'MOT17-09', 'MOT17-10',
                        'MOT17-11', 'MOT17-13']
elif MOT_set == 'MOT15':
    root = '/home/sean/data/MOT/2DMOT2015/train'
    train_video_list = ['TUD-Campus', 'ETH-Sunnyday', 'ETH-Pedcross2', 'ADL-Rundle-8', 'Venice-2',
                    'KITTI-17', 'TUD-Stadtmitte', 'ETH-Bahnhof', 'PETS09-S2L1', 'ADL-Rundle-6',
                    'KITTI-13']
write_file = 'gt_%s.txt'
truncate = False
for video in train_video_list:
    for line in open(os.path.join(root,video,'gt/gt.txt'), 'r'):
        if MOT_set == 'MOT17Det':
            frame_num, id, x, y, w, h, conf, cls, vis = line.split(',')
        elif MOT_set == 'MOT15':
            frame_num, id, x, y, w, h, conf, x, y, z = line.split(',')
        write_file_name = write_file % frame_num.zfill(6)
        write_file_path = os.path.join(root,video,'gt', write_file_name)
        if not os.path.exists(write_file_path):
            os.system(r'touch %s' % write_file_path)
        if truncate:
            with open(write_file_path, 'w') as wf:
                wf.truncate()
        elif MOT_set == 'MOT17Det' and int(conf) == 1 and float(vis) > 0.3:
            with open(write_file_path, 'a') as wf:
                wf.write(line)
        elif MOT_set == 'MOT15':
            with open(write_file_path, 'a') as wf:
                wf.write(line)